package com.edu.chat.web.integration.impl;

import com.edu.buservice.util.model.TenantUtil;
import com.edu.chat.web.config.InitMilvusConfig;
import com.edu.chat.web.dto.SearchArticleDTO;
import com.edu.chat.web.dto.SearchDocDTO;
import com.edu.chat.web.integration.MilvusClient;
import com.edu.chat.web.model.ArticleEntity;
import com.edu.chat.web.model.DocEntity;
import com.edu.robot.facade.VectorProcessingFacade;
import com.edu.robot.model.Request.TextEmbeddingRequest;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.vector.request.DeleteReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.DeleteResp;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.dubbo.config.annotation.DubboReference;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

@Component
@Slf4j
public class MilvusClientImpl implements MilvusClient {

	private static final Gson GSON = new Gson();
	@Autowired
	private MilvusClientV2 milvusClientV2;
	@Value("${milvus.search.score-threshold:0.6}")
	private float scoreThreshold;

	@DubboReference(interfaceClass = VectorProcessingFacade.class, timeout = 5000, check = false)
	private VectorProcessingFacade vectorProcessingFacade;


	@Override
	public List<SearchArticleDTO> search(String content,List<String> libraryIds) {
		return performSearch(content, libraryIds,
				InitMilvusConfig.COLLECTION_ARTICLE_NAME,
				Arrays.asList("id", "articleId", "title", "content", "libraryId", "tenantId"),
				this::processSearchResults
		);

	}

	@Override
	public List<SearchDocDTO> searchFile(String content, List<String> libraryId) {

		return performSearch(content, libraryId,
				InitMilvusConfig.COLLECTION_DOC_NAME,
				Arrays.asList("id", "documentId", "text", "libraryId", "tenantId"),
				this::processDocSearchResults
		);

	}

	@Override
	public void deleteFile(List<String> documentIds) {

		if (CollectionUtils.isEmpty(documentIds)) {
			log.warn("Delete document failed: documentIds is empty");
			return ;
		}


		final String filter = "tenantId == \"" + TenantUtil.getTenantId() + "\" AND documentId IN ["
				+ documentIds.stream()
				.map(Object::toString) // Convert each libraryId to String
				.collect(Collectors.joining(","))
				+ "]";

		milvusClientV2.delete(DeleteReq.builder()
				.collectionName(InitMilvusConfig.COLLECTION_DOC_NAME)
				.filter(filter)
				.build());

	}

	@Override
	public void deleteArticle(List<String> articleIds) {
		if (CollectionUtils.isEmpty(articleIds)) {
			log.warn("Delete article failed: articleIds is empty");
			return ;
		}

		final String filter = "tenantId == \"" + TenantUtil.getTenantId() + "\" AND articleId IN ["
				+ articleIds.stream()
				.map(Object::toString) // Convert each libraryId to String
				.collect(Collectors.joining(","))
				+ "]";

		milvusClientV2.delete(DeleteReq.builder()
				.collectionName(InitMilvusConfig.COLLECTION_ARTICLE_NAME)
				.filter(filter)
				.build());

	}

	@Override
	@Async("vectorizationThreadPool")
	public void insertData(List<ArticleEntity> articleList) {
		if (CollectionUtils.isEmpty(articleList)) {
			log.warn("Insert data failed: articleList is empty");
			return ;
		}

		final String tenantId = TenantUtil.getTenantId();

		List<JsonObject> jsonDataList = articleList.stream()
				.map(article ->{
					List<Float> floats = vectorProcessingFacade
							.vectorProcessing(new TextEmbeddingRequest(merge(article.getTitle(), article.getContent())))
							.stream()
							.map(Float::valueOf)
							.toList();
					JsonObject jsonObject = new JsonObject();
					jsonObject.addProperty("articleId", article.getArticleId());
					jsonObject.addProperty("title", article.getTitle());
					jsonObject.addProperty("content", article.getContent());
					jsonObject.addProperty("libraryId", article.getLibraryId());
					jsonObject.addProperty("tenantId", tenantId);
					jsonObject.add("vector",GSON.toJsonTree(floats));

					return jsonObject;
				})
				.toList();

		try {
			InsertReq insertReq = InsertReq.builder()
					.collectionName(InitMilvusConfig.COLLECTION_ARTICLE_NAME)
					.data(jsonDataList)
					.build();

			InsertResp insertResp = milvusClientV2.insert(insertReq);
			boolean success = insertResp.getInsertCnt() == articleList.size();
			if (!success) {
				log.warn("Partial data inserted, expected: {}, actual: {}",
						articleList.size(), insertResp.getInsertCnt());
			}
		} catch (Exception e) {
			log.error("Milvus data insertion failed", e);
		}
	}

	@Override
	@Async("vectorizationThreadPool")
	public void insertDocData(List<DocEntity> docList) {
		log.info("insertDocData doc:{}",docList);

		if (CollectionUtils.isEmpty(docList)) {
			log.warn("Insert data failed: docList is empty");
			return ;
		}

		final String tenantId = TenantUtil.getTenantId();

		List<JsonObject> jsonDataList = docList.stream()
				.map(doc ->{
					List<Float> floats = vectorProcessingFacade
							.vectorProcessing(new TextEmbeddingRequest(doc.getText()))
							.stream()
							.map(Float::valueOf)
							.toList();

					JsonObject jsonObject = new JsonObject();
					jsonObject.addProperty("documentId", doc.getDocumentId());
					jsonObject.addProperty("libraryId", doc.getLibraryId());
					jsonObject.addProperty("text", doc.getText());
					jsonObject.addProperty("tenantId", tenantId);
					jsonObject.add("vector",GSON.toJsonTree(floats));
					return jsonObject;
				})
				.toList();

		try {
			InsertReq insertReq = InsertReq.builder()
					.collectionName(InitMilvusConfig.COLLECTION_DOC_NAME)
					.data(jsonDataList)
					.build();
			milvusClientV2.insert(insertReq);

		}catch (Exception e){
			log.error("Milvus data insertion failed:{}", e.getMessage());
		}

	}

	private <T> List<T> performSearch(String content, List<String> libraryIds,
									  String collectionName, List<String> outputFields,
									  Function<SearchResp, List<T>> resultProcessor) {
		if (StringUtils.isBlank(content)) {
			log.warn("Search content is empty");
			return Collections.emptyList();
		}

		try {
			List<Float> vector = vectorProcessingFacade.vectorProcessing(new TextEmbeddingRequest(content))
					.stream()
					.map(Float::valueOf)
					.collect(Collectors.toList());
			SearchReq searchReq = buildSearchRequest(vector, libraryIds, collectionName, outputFields);
			SearchResp searchResp = milvusClientV2.search(searchReq);
			return resultProcessor.apply(searchResp);
		} catch (Exception e) {
			log.error("Search operation failed for content: {}", content, e);
			return Collections.emptyList();
		}
	}

	private SearchReq buildSearchRequest(List<Float> vector, List<String> libraryIds,
										 String collectionName, List<String> outputFields) {
		if (StringUtils.isBlank(TenantUtil.getTenantId())) {
			throw new IllegalArgumentException("Tenant ID must be provided");
		}
		if (CollectionUtils.isEmpty(libraryIds)) {
			throw new IllegalArgumentException("At least one library ID must be specified");
		}
		String filter = "tenantId == \"" + TenantUtil.getTenantId() + "\" AND libraryId IN ["
				+ libraryIds.stream()
				.map(Object::toString) // Convert each libraryId to String
				.collect(Collectors.joining(","))
				+ "]";

		return SearchReq.builder()
				.collectionName(collectionName)
				.data(Collections.singletonList(new FloatVec(vector)))
				.topK(InitMilvusConfig.topK)
				.outputFields(outputFields)
				.filter(filter)
				.build();
	}

	private List<SearchArticleDTO> processSearchResults(SearchResp searchResp) {
		return processGenericResults(searchResp, entity -> {
			SearchArticleDTO dto = new SearchArticleDTO();
			dto.setId(getLong(entity, "id"));
			dto.setArticleId(getString(entity, "articleId"));
			dto.setTitle(getString(entity, "title"));
			dto.setContent(getString(entity, "content"));
			dto.setLibraryId(getString(entity, "libraryId"));
			dto.setTenantId(getString(entity, "tenantId"));
			dto.setScore(getFloat(entity, "score"));
			return dto;
		});
	}

	private List<SearchDocDTO> processDocSearchResults(SearchResp searchResp) {
		return processGenericResults(searchResp, entity -> {
			SearchDocDTO dto = new SearchDocDTO();
			dto.setId(getLong(entity, "id"));
			dto.setDocumentId(getLong(entity, "documentId"));
			dto.setText(getString(entity, "text"));
			dto.setLibraryId(getString(entity, "libraryId"));
			dto.setTenantId(getString(entity, "tenantId"));
			dto.setScore(getFloat(entity, "score"));
			return dto;
		});
	}

	private <T> List<T> processGenericResults(SearchResp searchResp,
											  Function<Map<String, Object>, T> mapper) {
		return Optional.ofNullable(searchResp)
				.map(SearchResp::getSearchResults)
				.orElse(Collections.emptyList())
				.stream()
				.flatMap(List::stream)
				.filter(result -> result.getScore() >= scoreThreshold)
				.map(result -> {
					try {
						result.getEntity().put("score", result.getScore());
						return mapper.apply(result.getEntity());
					} catch (Exception e) {
						log.error("Result processing error", e);
						return null;
					}
				})
				.filter(Objects::nonNull)
				.collect(Collectors.toList());
	}

	private Long getLong(Map<String, Object> entity, String key) {
		return Optional.ofNullable(entity.get(key))
				.map(Object::toString)
				.map(Long::valueOf)
				.orElse(null);
	}

	private Float getFloat(Map<String, Object> entity, String key) {
		return Optional.ofNullable(entity.get(key))
				.map(Object::toString)
				.map(Float::valueOf)
				.orElse(null);
	}

	private String getString(Map<String, Object> entity, String key) {
		return Optional.ofNullable(entity.get(key))
				.map(Object::toString)
				.orElse(null);
	}

	public static String merge(String title, String content) {
		// 添加XML风格标记（需模型支持结构理解）
		return String.format(
				"<title>%s</title>\n<content>%s</content>",
				title, content
		);
	}

}